from cozepy import COZE_CN_BASE_URL, Coze, TokenAuth, Stream, WorkflowEvent, WorkflowEventType
from typing import Optional


def run_coze_workflow(user_input: str, workflow_id: str, api_token: str, base_url: str = COZE_CN_BASE_URL) -> Optional[str]:

    coze = Coze(auth=TokenAuth(token=api_token), base_url=base_url)
    final_message = ""

    def handle_workflow_iterator(stream: Stream[WorkflowEvent]):
        nonlocal final_message
        for event in stream:
            if event.event == WorkflowEventType.MESSAGE:
                if event.message and event.message.content:
                    final_message += event.message.content
            elif event.event == WorkflowEventType.ERROR:
                print("发生错误:", event.error)
                return None
            elif event.event == WorkflowEventType.INTERRUPT:
                handle_workflow_iterator(
                    coze.workflows.runs.resume(
                        workflow_id=workflow_id,
                        event_id=event.interrupt.interrupt_data.event_id,
                        resume_data="继续",
                        interrupt_type=event.interrupt.interrupt_data.type,
                    )
                )

    # 启动工作流
    stream = coze.workflows.runs.stream(
        workflow_id=workflow_id,
        parameters={"input": user_input}
    )

    handle_workflow_iterator(stream)

    return final_message.strip() if final_message else None

# 使用示例
if __name__ == "__main__":
    # 配置参数
    # WORKFLOW_ID = "7528679746101182464"
    WORKFLOW_ID="7517921086106157119"
    # API_TOKEN = "pat_ljYM5e6wZTrmM8tdzI3BQeMnrLutfRGGwBAxqq7X4xjs4zCqUFQABiwvbNsOtpip"
    API_TOKEN="pat_nnTFo3UOS0uEcaIOn4aY7GPHg1GjnFQgvX6WHtefmQW2jhnk7z6nt2tq2wbGFFws"
    # 运行工作流
    a=run_coze_workflow("你",workflow_id=WORKFLOW_ID, api_token=API_TOKEN)
    print(a)